package com.qf.util;

import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶 - 每个对象代表一个令牌桶
 *
 * 对象 -> lua脚本 -> redis服务
 *
 * 事件总线 - N个绑定对象手动注册到Spring中
 *
 * ip限流
 * 用户限流
 * url限流.....
 *
 */
public class TokenTong {

    //在redis中，令牌桶的key
    private String key;
    //令牌的最大容量
    private Integer maxTokens;
    //令牌的生成速率 - 每秒生成多少令牌
    private Integer secTokens;
    //操作redis的对象
    private StringRedisTemplate redisTemplate;

    //初始化令牌桶的lua脚本
    private String initLua = "--初始化令牌桶的key\n" +
            "local key = 'tokentong_'..KEYS[1]\n" +
            "--令牌桶最大容量\n" +
            "local maxTokens = ARGV[1]\n" +
            "--每秒生成的速率\n" +
            "local secTokens = ARGV[2]\n" +
            "--获得当前时间\n" +
            "local timeArray = redis.call('time')\n" +
            "--当前时间的微秒值\n" +
            "local nowTime = timeArray[1] * 1000000 + timeArray[2]\n" +
            "\n" +
            "--判断令牌桶是否存在\n" +
            "local flag = redis.call('exists', key)\n" +
            "if flag == 0 then\n" +
            "\t--初始化令牌桶\n" +
            "\tredis.call('hmset', key, 'hasTokens', maxTokens, 'maxTokens', maxTokens, 'secTokens', secTokens, 'nextTime', nowTime)\n" +
            "\treturn 1\n" +
            "end\n" +
            "--未初始化\n" +
            "return 0";

    //获取令牌的lua脚本
    private String getTokenLua = "--获得令牌桶key\n" +
            "local key = 'tokentong_'..KEYS[1]\n" +
            "--获得需要获取的令牌数量\n" +
            "local getTokens = ARGV[1]\n" +
            "--获得令牌桶的剩余令牌数\n" +
            "local hasTokens = tonumber(redis.call('hget', key, 'hasTokens'))\n" +
            "--获得令牌桶的最大容量\n" +
            "local maxTokens = tonumber(redis.call('hget', key, 'maxTokens'))\n" +
            "--获得令牌桶的生成速率 - 每秒生成的个数\n" +
            "local secTokens = tonumber(redis.call('hget', key, 'secTokens'))\n" +
            "--下一次可以生成令牌的时间\n" +
            "local nextTimes =  tonumber(redis.call('hget', key, 'nextTime'))\n" +
            "--允许的超时时间\n" +
            "local timeout = tonumber(ARGV[2] or -1)\n" +
            "\n" +
            "--当前时间\n" +
            "local timeArray = redis.call('time')\n" +
            "--当前时间的微秒值\n" +
            "local nowTimes = timeArray[1] * 1000000 + timeArray[2]\n" +
            "\n" +
            "--单个令牌的生成耗时\n" +
            "local singleTokenTime = 1000000 / secTokens\n" +
            "\n" +
            "--判断是否有超时时间\n" +
            "if timeout ~= -1 then\n" +
            "  --判断等待时间是否会超过超时时间，如果会超过，则直接返回-1  \n" +
            "  if timeout < nextTimes - nowTimes then\n" +
            "    --返回-1，表示在超时时间范围内，无法获取令牌\n" +
            "    return -1\n" +
            "  end\n" +
            "end\n" +
            "\n" +
            "--计算是否有新的令牌生成\n" +
            "if nowTimes > nextTimes then\n" +
            "  --计算两个时间的时间差\n" +
            "  local xTime = nowTimes - nextTimes\n" +
            "  --计算会生成的令牌数\n" +
            "  local createTokens = xTime / singleTokenTime\n" +
            "  --修改参数\n" +
            "  hasTokens = math.min(hasTokens + createTokens, maxTokens)\n" +
            "  nextTimes = nowTimes\n" +
            "end \n" +
            "\n" +
            "--获取令牌\n" +
            "--计算能够获取的令牌\n" +
            "local canGetTokens = math.min(getTokens, hasTokens)\n" +
            "--计算需要预支的令牌\n" +
            "local yuzhiTokens = getTokens - canGetTokens\n" +
            "--计算剩余的令牌\n" +
            "hasTokens = hasTokens - canGetTokens\n" +
            "--计算预支令牌后，所需要时间\n" +
            "local yuzhiTimes = singleTokenTime * yuzhiTokens\n" +
            "--重新设置参数 - 修改redis\n" +
            "redis.call('hmset', key, 'hasTokens', hasTokens, 'nextTime', nextTimes + yuzhiTimes)\n" +
            "--返回需要等待的时间 - 微秒\n" +
            "return nextTimes - nowTimes\n";

    /**
     * 令牌桶初始化
     */
    public TokenTong(String key, Integer maxTokens, Integer secTokens){
        this.key = key;
        this.maxTokens = maxTokens;
        this.secTokens = secTokens;
        this.redisTemplate = SpringUtil.getBean(StringRedisTemplate.class);

        //在redis中初始化令牌桶
        long result = redisTemplate.execute(new DefaultRedisScript<>(initLua, Long.class),
                Collections.singletonList(key),
                maxTokens + "", secTokens + "");
    }

    /**
     * 领取令牌的方法
     * 1、直接领取令牌 - 范围需要等待的时间
     * 2、直接领取令牌，并且设置超时时间 - 如果需要等待的时间，超过了指定的超时时间，则直接报错
     * 3、直接领取令牌，如果不能立刻获得令牌，直接报错
     */
    //1、直接领取令牌 - 返回需要等待的时间
    public long getTokens(int tokens){

        Long wait = redisTemplate.execute(new DefaultRedisScript<Long>(getTokenLua, Long.class),
                Collections.singletonList(key),
                tokens + "");

        if (wait > 0) {
            //将微秒转毫秒
            long ms = TimeUnit.MICROSECONDS.toMillis(wait);
            try {
                Thread.sleep(ms);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }

        return wait;
    }

    //2、直接领取令牌 - 返回需要等待的时间
    public boolean getTokens(int tokens, long timeout, TimeUnit unit){

        Long wait = redisTemplate.execute(new DefaultRedisScript<Long>(getTokenLua, Long.class),
                Collections.singletonList(key),
                tokens + "", unit.toMicros(timeout) + "");

        if (wait == -1) {
            return false;
        } else if (wait > 0){
            //将微秒转毫秒
            long ms = TimeUnit.MICROSECONDS.toMillis(wait);
            try {
                Thread.sleep(ms);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        return true;
    }

    //3、直接领取令牌，如果不能立刻获得令牌，直接报错
    public boolean getTokensNow(int tokens){
        return getTokens(tokens, 0, TimeUnit.MICROSECONDS);
    }
}
